import sounddevice as sd
import numpy as np
import queue
from collections import deque
import threading
import time
import asyncio
from io import BytesIO

from text_process import process_text
from baidu_asr_client import baidu_asr_send

# 尝试导入TTS相关库
try:
    import edge_tts
    import pydub
    from pydub import AudioSegment

    EDGE_TTS_AVAILABLE = True
except ImportError:
    EDGE_TTS_AVAILABLE = False

VAD_THRESHOLD = 2000
VAD_SILENCE_FRAMES = 20
PRE_RECORDING_FRAMES = 5

# --- 音频流参数 ---
SAMPLERATE = 16000
CHANNELS = 1
DTYPE = 'int16'
BLOCK_SIZE = 1024

# 用于线程间通信的队列
audio_queue = queue.Queue()
response_queue = queue.Queue()  # 文本处理结果队列
playback_queue = queue.Queue()  # 音频播放队列
playback_finished_event = threading.Event()  # 播放完成事件

# --- VAD 状态变量 ---
is_recording = False
silence_counter = 0
recording_buffer = []
pre_recording_queue = deque(maxlen=PRE_RECORDING_FRAMES)

# --- 播放控制变量 ---
is_playing = False  # 是否正在播放音频
should_pause_recording = False  # 是否应该暂停录音
audio_stream = None  # 音频流对象

# --- TTS 配置 ---
TTS_VOICE = "zh-CN-YunjianNeural"  # Edge TTS 中文女声
TTS_MODE = "edge"  # 可选: "edge", "pyttsx3", "beep"


# 其他可用选项:
# "zh-CN-YunxiNeural"     # 中文男声
# "zh-CN-XiaoyiNeural"    # 中文女声（温柔）
# "zh-CN-YunjianNeural"   # 中文男声（成熟）

def audio_callback(indata, frames, time, status):
    global is_recording, silence_counter, recording_buffer, pre_recording_queue, should_pause_recording

    if status:
        print(status)

    # 如果正在播放音频，暂停录音处理
    if should_pause_recording:
        return

    # 计算当前音频区块的能量 (RMS)
    # 注意：indata是int16类型，计算时最好转为浮点数防止溢出
    rms = np.sqrt(np.mean(indata.astype(np.float32) ** 2))

    # 可以在这里取消注释来调试阈值
    # print(f"Current RMS: {rms:.2f}")

    if is_recording:
        # --- 处于录音状态 ---
        recording_buffer.append(indata.copy())
        if rms < VAD_THRESHOLD:
            silence_counter += 1
            if silence_counter > VAD_SILENCE_FRAMES:
                print("检测到静音结束，完成录制。")
                # 将录音数据从缓冲区拼接成一个完整的Numpy数组
                full_recording = np.concatenate(recording_buffer)
                # 将完整的录音放入队列，供其他线程处理
                audio_queue.put(full_recording)

                # 重置状态
                is_recording = False
                silence_counter = 0
                recording_buffer.clear()
        else:
            # 有语音活动，重置静音计数器
            silence_counter = 0
    else:
        # --- 处于等待状态 ---
        # 无论如何都将当前区块加入 pre_recording_queue
        # deque 会自动管理长度
        pre_recording_queue.append(indata.copy())

        if rms > VAD_THRESHOLD:
            print("检测到语音，开始录制...")
            is_recording = True
            # 将触发前的几个区块也加入缓冲区
            recording_buffer.extend(list(pre_recording_queue))
            silence_counter = 0


def start_audio_stream():
    """启动音频录制流"""
    global audio_stream
    try:
        if audio_stream is not None:
            audio_stream.close()

        audio_stream = sd.InputStream(
            samplerate=SAMPLERATE,
            channels=CHANNELS,
            dtype=DTYPE,
            blocksize=BLOCK_SIZE,
            callback=audio_callback
        )
        audio_stream.start()
        print("[音频流] 录音流已启动")
    except Exception as e:
        print(f"[音频流] 启动录音流失败: {e}")


def stop_audio_stream():
    """停止音频录制流"""
    global audio_stream, is_recording, silence_counter, recording_buffer, pre_recording_queue
    try:
        if audio_stream is not None:
            audio_stream.stop()
            audio_stream.close()
            audio_stream = None

        # 重置录音状态
        is_recording = False
        silence_counter = 0
        recording_buffer.clear()
        pre_recording_queue.clear()

        print("[音频流] 录音流已停止，状态已重置")
    except Exception as e:
        print(f"[音频流] 停止录音流失败: {e}")

def tts_worker():
    """
    TTS工作线程：
    - 从 response_queue 获取文本
    - 使用edge-tts合成语音流
    - 将音频数据块放入 playback_queue
    """
    print("[TTS线程] 启动，准备合成语音...")
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    async def do_edge_tts(text):
        print(f"[TTS线程] 正在合成: {text}")
        try:
            communicate = edge_tts.Communicate(text, TTS_VOICE)
            async for chunk in communicate.stream():
                if chunk["type"] == "audio":
                    playback_queue.put(chunk["data"])
            playback_queue.put(None)  # 标记一个文本的音频流结束
            print("[TTS线程] 语音合成完成")
        except Exception as e:
            print(f"[TTS线程] Edge TTS合成失败: {e}")
            playback_queue.put(None)

    def do_beep_sound():
        print("[TTS线程] 生成提示音")
        duration = 2.0
        frequency = 440
        sample_rate = 44100
        t = np.linspace(0, duration, int(sample_rate * duration), False)
        audio_data = np.sin(frequency * 2 * np.pi * t) * 0.3

        # 转换为16位PCM
        audio_16bit = (audio_data * 32767).astype(np.int16)
        playback_queue.put(audio_16bit.tobytes())
        playback_queue.put(None)

    while True:
        text_to_speak = response_queue.get()
        if text_to_speak is None:
            break

        if TTS_MODE == "edge" and EDGE_TTS_AVAILABLE:
            loop.run_until_complete(do_edge_tts(text_to_speak))
        # elif TTS_MODE == "pyttsx3" and PYTTSX3_AVAILABLE:
        #     do_pyttsx3_tts(text_to_speak)
        else:
            do_beep_sound()

    loop.close()


def playback_worker():
    """
    播放工作线程：
    - 从 playback_queue 获取音频数据块
    - 使用sounddevice流式播放
    """
    print("[播放线程] 启动，等待音频流...")

    while True:
        # 收集一个完整语音的所有音频块
        full_audio_data = bytearray()

        while True:
            audio_chunk = playback_queue.get()

            if audio_chunk is None:  # 一个文本播放完毕的标记
                break

            full_audio_data.extend(audio_chunk)

        if not full_audio_data:
            print("[播放线程] 没有收到音频数据，跳过")
            continue

        try:
            if TTS_MODE == "edge" and EDGE_TTS_AVAILABLE:
                # Edge TTS输出MP3格式，需要解码
                print("[播放线程] 正在解码MP3音频...")
                mp3_stream = BytesIO(full_audio_data)
                audio_segment = AudioSegment.from_file(mp3_stream, format="mp3")
                pcm_data = np.array(audio_segment.get_array_of_samples(), dtype=np.int16)

                print(f"[播放线程] 正在播放合成的语音... (采样率: {audio_segment.frame_rate}Hz, 长度: {len(pcm_data)})")
                sd.play(pcm_data, samplerate=audio_segment.frame_rate)
                sd.wait()

            else:
                # 提示音是PCM格式
                pcm_data = np.frombuffer(full_audio_data, dtype=np.int16)
                print(f"[播放线程] 正在播放提示音...")
                sd.play(pcm_data, samplerate=44100)
                sd.wait()

            print("[播放线程] 播放完毕")

            # 设置播放完成事件
            playback_finished_event.set()

        except Exception as e:
            print(f"[播放线程] 播放失败: {e}")
            import traceback
            traceback.print_exc()

            # 即使播放失败也要设置完成事件
            playback_finished_event.set()


def process_and_play_worker(text: str):
    """
    工作线程函数：处理文本并启动TTS播放流程
    """
    global should_pause_recording, is_playing

    try:
        # 停止录音流避免回音
        print("[主处理] 停止录音流，开始处理...")
        stop_audio_stream()

        should_pause_recording = True
        is_playing = True

        # 清除播放完成事件
        playback_finished_event.clear()

        # 处理文本
        processed_text = process_text(text)

        # 将处理结果发送到TTS队列
        print("[主处理] 发送文本到TTS队列...")
        response_queue.put(processed_text)

        # 等待播放完成事件
        print("[主处理] 等待播放完成...")
        playback_finished = playback_finished_event.wait(timeout=30)  # 最多等待30秒

        if playback_finished:
            print("[主处理] 播放已完成")
        else:
            print("[主处理] 播放等待超时")

        # 播放完成后等待一点时间，让音频完全结束
        print("[主处理] 等待音频完全结束...")
        time.sleep(1.5)  # 给音频输出设备时间完全停止

    except Exception as e:
        print(f"[主处理] 处理过程中出现错误: {e}")
        import traceback
        traceback.print_exc()

    finally:
        # 恢复录音流
        is_playing = False
        should_pause_recording = False
        print("[主处理] 重新启动录音流...")
        start_audio_stream()

        # 额外等待一点时间让录音流稳定
        time.sleep(0.5)



# 主程序
if __name__ == "__main__":
    # 检查可用的TTS引擎并设置默认模式
    if EDGE_TTS_AVAILABLE:
        TTS_MODE = "edge"
        print("✓ 检测到 Edge TTS 支持")
    else:
        TTS_MODE = "beep"
        print("⚠ 未检测到TTS支持，将使用提示音")

    print("\n=== 数字人语音助手 ===")
    print(f"TTS模式: {TTS_MODE}")
    print(f"语音模型: {TTS_VOICE if TTS_MODE == 'edge' else 'N/A'}")

    tts_thread = threading.Thread(target=tts_worker, daemon=True)
    tts_thread.start()

    playback_thread = threading.Thread(target=playback_worker, daemon=True)
    playback_thread.start()

    start_audio_stream()

    print("正在监听麦克风... 请说话。")

    try:
        while True:
            recorded_audio = audio_queue.get()

            result = baidu_asr_send(recorded_audio, SAMPLERATE)

            if result is not None and result.strip():
                print(f"识别结果: {result}")

                # 在新线程中处理文本，避免阻塞主线程
                worker_thread = threading.Thread(
                    target=process_and_play_worker,
                    args=(result,),
                    daemon=True
                )
                worker_thread.start()

            elif result is not None:
                print("识别结果为空，继续监听...")
            else:
                print("语音识别失败，继续监听...")

    except KeyboardInterrupt:
        print("\n正在关闭程序...")

        # 停止音频流
        stop_audio_stream()

        # 停止TTS线程
        response_queue.put(None)

        print("程序已退出。")